import os
from PIL import Image, ImageOps
import numpy as np
from scipy.ndimage import label
from concurrent.futures import ThreadPoolExecutor, as_completed
from tqdm import tqdm

def split_and_crop_image(image_path, output_dir):
    try:
        # 打开图片并转换为RGBA模式
        image = Image.open(image_path).convert("RGBA")
        data = np.array(image)

        # 提取Alpha通道
        alpha_channel = data[:, :, 3]

        # 找到非透明区域
        binary_mask = (alpha_channel > 0).astype(np.uint8)
        labeled_array, num_features = label(binary_mask)

        if num_features == 0:
            return []  # 如果没有非透明区域，返回空列表

        # 找到最靠近下边的连通区域
        bottom_region_label = None
        bottom_y = -1
        for label_value in range(1, num_features + 1):
            region_mask = (labeled_array == label_value)
            region_coords = np.argwhere(region_mask)
            region_bottom_y = np.max(region_coords[:, 0])
            if region_bottom_y > bottom_y:
                bottom_y = region_bottom_y
                bottom_region_label = label_value

        if bottom_region_label is None:
            return []  # 如果未找到底部图案，返回空列表

        # 创建掩码：最靠近下边的图案
        bottom_mask = (labeled_array == bottom_region_label).astype(np.uint8) * 255

        # 创建剩余部分：将底部图案区域的Alpha通道设为0
        remaining_alpha = alpha_channel.copy()
        remaining_alpha[bottom_mask > 0] = 0  # 将底部图案区域设为透明
        remaining_data = data.copy()
        remaining_data[:, :, 3] = remaining_alpha

        # 转换为PIL图像
        bottom_image = Image.fromarray(data * (bottom_mask // 255)[:, :, np.newaxis], mode="RGBA")
        remaining_image = Image.fromarray(remaining_data, mode="RGBA")

        # 裁剪到有像素的内容贴着四周边缘
        def crop_to_content(image):
            bbox = image.getbbox()
            return image.crop(bbox) if bbox else image  # 如果无有效区域，返回原图

        bottom_image = crop_to_content(bottom_image)
        remaining_image = crop_to_content(remaining_image)

        # 保存结果
        base_name = os.path.splitext(os.path.basename(image_path))[0]
        bottom_output = os.path.join(output_dir, f"{base_name}_bottom.png")
        remaining_output = os.path.join(output_dir, f"{base_name}_remaining.png")

        bottom_image.save(bottom_output)
        remaining_image.save(remaining_output)

        return [bottom_output, remaining_output]
    except Exception as e:
        print(f"Error processing {image_path}: {str(e)}")
        return []

def batch_process_images(input_dir, output_dir, max_workers=8):
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)

    # 收集所有图片路径
    image_list = [
        os.path.join(input_dir, f)
        for f in os.listdir(input_dir)
        if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp'))
    ]

    # 使用多线程处理
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = [
            executor.submit(split_and_crop_image, img_path, output_dir)
            for img_path in image_list
        ]

        # 显示进度条
        for future in tqdm(as_completed(futures), total=len(futures), desc="Processing Images"):
            try:
                _ = future.result()  # 获取结果（可选）
            except Exception as e:
                print(f"Error in processing thread: {str(e)}")


# 示例调用
if __name__ == "__main__":
    input_directory = "D:/Desk/0312-炳哥尺寸转换/尺寸删除"
    output_directory = "D:/Desk/0312-炳哥尺寸转换/result"
    batch_process_images(input_directory, output_directory, max_workers=6)